package com.xteam.xnetworkdisk.websocket;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

@Component
public class WebSocketSessionManager {

    private final Map<String, WebSocketSession> sessionMap = new HashMap<>();

    public void addSession(String md5, WebSocketSession session) {
        sessionMap.put(md5, session);
    }

    public void removeSession(String md5) {
        sessionMap.remove(md5);
    }

    public void sendProgressUpdate(String md5, String message) throws IOException {
        WebSocketSession session = sessionMap.get(md5);
        if (session != null && session.isOpen()) {
            session.sendMessage(new TextMessage(message));
        }
    }
    public void closeSession(String md5) throws IOException {
        WebSocketSession session = sessionMap.get(md5);
        if (session != null && session.isOpen()) {
            session.close(CloseStatus.NORMAL);
            sessionMap.remove(md5);
        }
    }
}
